#include <string>
#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <vector>
#include <sstream>
#include <fstream>
#include <boost/lexical_cast.hpp>
#define DEBUG_LEVEL 6
#include "blast_results.h"
#include "batch_job.h"
#include "portable_timer.h"
#include "debug_func.h"
#include "codeml.h"
#include "perlutil.h"
#include "progress_indicator.h"
#include "pg_db_cursor.h"
#include "pg_copy_guard.h"
#include "ultoa.h"
#define _MY_STL_MACROS
#include "stlhelper.h"
#include "auto_restore_dir.h"
#include "string_hash.h"
#include "get_arg.h"
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <ext/hash_map>
#include <boost/array.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
using std::cerr;
using std::cout;
using std::string;
using std::vector;
using std::sort;
using std::replace;
using boost::lexical_cast;
using boost::filesystem::path;
using std::ostringstream;
using std::istringstream;
using std::ofstream;
using std::ostream;
using std::ifstream;


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	KaKs

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

	
	
	
	
	
	









struct t_pair_data
{
	string			pi1;
	string			pi2;
	unsigned		cdna_index1;
	unsigned		cdna_index2;
	std::string		gap_str1;
	std::string		gap_str2;
	unsigned 		align_beg;
	unsigned 		align_len;
	t_pair_data(unsigned a1, unsigned a2, const string& g1, const string& g2,
				const string& p1, const string& p2,
				const string& s1, const string& s2, unsigned b, unsigned l):
				pi1(g1 + "\t" + p1), pi2(g2 + "\t" + p2),
				cdna_index1(a1),
				cdna_index2(a2),
				gap_str1(s1), gap_str2(s2), align_beg(b), align_len(l){}
};

typedef vector<t_pair_data> t_vec_pair_data;


//________________________________________________________________________________________

//		get_cdna_sequences

//________________________________________________________________________________________
// hash map
typedef __gnu_cxx::hash_map<string, unsigned, hash_string> t_map_acc_code_index;

bool get_cdna_sequences(t_pg_cursor&			pg_db,
						const vector<string>&	acc_codes,
						vector<string>&			cdna_sequences,
						const string&			batch,
						t_program_args& 		args)
{
	debug_func_cerr(5);
	
	
	//
	// create cursor
	//
	string sql_cmd = "SELECT "
						"distinct on (c.acc_code) c.acc_code, c.cdna "
					"FROM "
						"cdna_sequences c, "
						"ens_id e, "	
						"("
							"SELECT "
								"gene_id1 AS gene_id "
							"FROM "
								"batch_gene_pairs "
							"WHERE "
								"batch = " + batch + " AND "
								"category = '" + args.job_category + "' "
							"UNION "
							"SELECT "
								"gene_id2 AS gene_id "
							"FROM "
								"batch_gene_pairs "
							"WHERE "
								"batch = " + batch + " AND "
								"category = '" + args.job_category + "' "
						") as a "
					"WHERE "
						"c.acc_code = e.acc_code AND "
						"e.gene_id = a.gene_id";

	if (!pg_db.Declare(sql_cmd, false))
		return pg_print_sql_err("create a cursor to retrieve cdna sequences",
								pg_db,
								sql_cmd, cerr);
	//
	// map accession codes to indices so that cdna_sequences can be stored in order
	//
	t_map_acc_code_index map_acc_code_index;
	for (unsigned i = 0; i < acc_codes.size(); ++i)
		map_acc_code_index[acc_codes[i]] = i;
	cdna_sequences.assign(acc_codes.size(),  string());
	
	
	//
	// get data in batches
	//
	unsigned cnt_retrieved = 0;
	unsigned batch_size = 50;
	unsigned index = 0;
	while (pg_db.Fetch(batch_size))
	{
		int curr_batch_size = pg_db.Tuples();
		cnt_retrieved += curr_batch_size;

		if (!curr_batch_size)
			break;
		for (int i = 0; i < curr_batch_size; ++i)
		{
			t_map_acc_code_index::iterator iter = map_acc_code_index.find(string(pg_db.getvalue(i, 0)));
			if (iter != map_acc_code_index.end())
				cdna_sequences[iter->second] = pg_db.getvalue(i, 1);
			// else is not present in alignment
		}
	}

	if (!pg_db.Close())
		return pg_print_sql_err("close the cursor to retrieve cdna sequences ",
								pg_db, "CLOSE CURSOR",
								cerr);
	if (cnt_retrieved < acc_codes.size())
		throw std::logic_error("Too few cdna retrieved\n");
	
	for (int i = 0; i < cdna_sequences.size(); ++i)
		if (cdna_sequences[i].length() == 0)
			args.CERR() << "ERROR\nCould not retrieve cdna for "
						<< acc_codes[i] << "=" << i << "\n";

	return true;
}

//________________________________________________________________________________________

//		expand_gaps_in_cdna

//________________________________________________________________________________________
bool expand_gaps_in_cdna(const string& gaps_str, string& cdna)
{
	debug_func_cerr(10);
	// empty string encodes no gaps!

	if (! gaps_str.length())
		return true;

	std::vector<t_gap> gaps;
	if (!gap_str_to_array(gaps_str, gaps))
		return false;

	return expand_gaps_cdna(gaps, cdna);
}

//________________________________________________________________________________________

//		print_on_error

//			for debugging

//________________________________________________________________________________________
void print_on_error(t_program_args& args,
					const string& orig_sequence1,
					const string& orig_sequence2,
					const string& sequence1,
					const string& sequence2,
					const string& gap_str1, const string& gap_str2,
					const string& name1,
					const string& name2,
					unsigned align_len)
{
	args.CERR() << "names1 = " << name1.length() << "\t" << name1 << "\n";
	args.CERR() << "names2 = " << name2.length() << "\t" << name2 << "\n";
	args.CERR() << "orig_sequence1 = " << orig_sequence1.length() << "\t" << orig_sequence1 << "\n";
	args.CERR() << "orig_sequence2 = " << orig_sequence2.length() << "\t" << orig_sequence2 << "\n";
	args.CERR() << "gap_str1 = " << gap_str1 << "\n";
	args.CERR() << "gap_str2 = " << gap_str2 << "\n";
	args.CERR() << "sequence1 = " << sequence1.length() << "\t" << sequence1 << "\n";
	args.CERR() << "sequence2 = " << sequence2.length() << "\t" << sequence2 << "\n";
	args.CERR() << "align len = " << align_len * 3 << "\n\n";
	args.CERR().flush();
}



//________________________________________________________________________________________

//		insert_gaps_into_sequence_pair

//________________________________________________________________________________________
bool insert_gaps_into_sequence_pair(	const t_pair_data&			pair_data,
										string&						sequence1,
										string&						sequence2,	
										string&						error_str)
{
	try
	{

		if (!pair_data.align_len)
		{
			error_str = "Zero length alignment!\n";
			return false;
		}


		if (!sequence1.length() || !sequence2.length())
		{
			error_str = "Sequences missing.";
			return false;
		}

		if (!expand_gaps_in_cdna(pair_data.gap_str1, sequence1)	||
			!expand_gaps_in_cdna(pair_data.gap_str2, sequence2))
		{
			error_str = "Could not expand gaps.";
			return false;
		}
		if (!sequence1.length() || !sequence2.length())
		{
			error_str = "Sequence lengths mismatch after inserting gaps.";
			return false;
		}
		if (pair_data.align_beg * 3 > sequence1.length())
		{
			error_str = "Align start too large.";
			return false;
		}

		// remove non-aligned region
		sequence1.erase(0, pair_data.align_beg * 3);
		sequence2.erase(0, pair_data.align_beg * 3);
		
		if (pair_data.align_len * 3 > sequence1.length())
		{
			error_str = "Align length too large.";
			return false;
		}
		
		sequence1.erase(pair_data.align_len * 3);
		sequence2.erase(pair_data.align_len * 3);
		return true;

	}
	catch (...)
	{
		error_str = "Unknown gap insertion error\n";
		return false;
	}

}

//________________________________________________________________________________________

//		get_kaks_results

//________________________________________________________________________________________
void get_kaks_results(  const t_vec_pair_data&		pair_data,
						const vector<string>&		acc_codes,
						const vector<string>&		cdna_sequences,
						vector<t_codeml_results>&	results,
						t_progress_indicator& 		dots,
						t_program_args&				args)
{
	dots.print_indicator("KaKs:");
	debug_func_cerr(5);
	for (unsigned i = 0; i < pair_data.size(); ++i, ++dots)
	{
		string error_str;
		string seq1	= cdna_sequences[pair_data[i].cdna_index1];
		string seq2 = cdna_sequences[pair_data[i].cdna_index2];
		if (insert_gaps_into_sequence_pair(pair_data[i], seq1, seq2, error_str))
		{
			kaks_analyse(seq1, seq2, i, i, results, args.paml + " kaks.ctl", args.codeml);
			continue;
		}
		
		
		//	
		//	Error:
		//
		// save to db
		results.push_back(t_codeml_results(i, i, error_str));
		
		//print error
		args.CERR() << "ERROR: " << error_str << "\n";

		// print original ungapped sequences
		//
		print_on_error(args,
						cdna_sequences[pair_data[i].cdna_index1],
						cdna_sequences[pair_data[i].cdna_index2],
						seq1,
						seq2,
						pair_data[i].gap_str1, pair_data[i].gap_str2,
						pair_data[i].pi1 + "\t" + acc_codes[pair_data[i].cdna_index1],
						pair_data[i].pi2 + "\t" + acc_codes[pair_data[i].cdna_index2],
						pair_data[i].align_len);
		dots.print_indicator('!');
		
	}
}

//________________________________________________________________________________________

//		write_kaks_results

//________________________________________________________________________________________
ostream& print_sequence_descriptors(ostream&						os,
									unsigned						index,
									const t_vec_pair_data& 			pair_data,
									const vector<string>& 			acc_codes)
{
	os << pair_data[index].pi1 + '\t' + acc_codes[pair_data[index].cdna_index1] + '\t' +
			pair_data[index].pi2 + '\t' + acc_codes[pair_data[index].cdna_index2] + '\t';
	return os;
}

bool write_kaks_results(t_pg_cursor& 						pg_db, 		
							const t_vec_pair_data& 			pair_data,
							const vector<string>& 			acc_codes,
							const vector<t_codeml_results>& results,
							t_program_args&					args)
{
	debug_func_cerr(5);
	{
		t_pg_copy_guard copy_buffer(pg_db, args.output_table +
											"(gene_id1, pi1, acc_code1, "
											" gene_id2, pi2, acc_code2, "
											+ t_codeml_results::comma_delimited_fields()
											+ ")",
									cerr);
		if (!copy_buffer)
			return false;
		ostringstream ostr;
		for (unsigned i = 0; i < results.size(); ++i)
			print_sequence_descriptors(ostr, results[i].seq_tag1, pair_data, acc_codes)
			<<	results[i]				<<	'\n';
		copy_buffer += ostr.str();
		if (!copy_buffer.end_copy())
		{
			std::cout << copy_buffer.buf;
			return false;
		}
	}
	return true;
}



//________________________________________________________________________________________

//		t_acc_code_index

//________________________________________________________________________________________
struct t_acc_code_index
{
	string		acc_code;
	unsigned	index;
	t_acc_code_index(const string& s, unsigned i) : acc_code(s), index(i) {}
	bool operator < (const t_acc_code_index& o) const { return acc_code < o.acc_code; }
};



//________________________________________________________________________________________

//		construct_unique_acc_code_list

//________________________________________________________________________________________
void construct_unique_acc_code_list(vector<t_acc_code_index>&	vec_acc_code_indices,
									t_vec_pair_data&			pair_data,
									vector<string>&				acc_codes)
{
	debug_func_cerr(5);
	vector<unsigned> acc_code_index_map(vec_acc_code_indices.size());
	sort(beg_to_end(vec_acc_code_indices));
	// add first accession code
	acc_code_index_map[vec_acc_code_indices[0].index] = 0;
	acc_codes.push_back(vec_acc_code_indices[0].acc_code);

	// add successive acc_codes if different from previous
	for (unsigned i = 1, unique_i = 0; i < vec_acc_code_indices.size(); ++i)
	{
		if (vec_acc_code_indices[i].acc_code != vec_acc_code_indices[i-1].acc_code)
		{
			++unique_i;
			acc_codes.push_back(vec_acc_code_indices[i].acc_code);
		}
		acc_code_index_map[vec_acc_code_indices[i].index] = unique_i;
	}
	// map t_pair_data offsets back into the list of accession codes
	for (unsigned i = 0; i < pair_data.size(); ++i)
	{
		pair_data[i].cdna_index1 = acc_code_index_map[i * 2];
		pair_data[i].cdna_index2 = acc_code_index_map[i * 2 + 1];
	}
}

//________________________________________________________________________________________

//		get_sequence_pair_data

//________________________________________________________________________________________
bool get_sequence_pair_data(t_pg_cursor&		pg_db,
							t_program_args&		args,
							const string& 		batch_num,
							t_vec_pair_data&	pair_data,
							vector<string>&		acc_codes)
{
	debug_func_cerr(5);

	// create cursor
	string sql_cmd ="SELECT distinct on (gene_id1, gene_id2) "
						"gene_id1, gene_id2, pi1, pi2, acc_code1, acc_code2, "
							"aln_gaps1, aln_gaps2, bl2seq_align_start, bl2seq_align_len "
					"FROM " +
						args.input_table + " JOIN "
						"batch_gene_pairs USING (gene_id1, gene_id2) "
					"WHERE "
						"batch = " + batch_num + 					" AND "
						"category = '" + args.job_category + "'";
	if (!pg_db.Declare(sql_cmd, false))
		return pg_print_sql_err("create a cursor to retrieve batch " + batch_num,
								pg_db,
								sql_cmd, cerr);

	
	// get data in batches
	unsigned batch_size = 50;
	unsigned index = 0;
	vector<t_acc_code_index>	vec_acc_code_indices;
	while (pg_db.Fetch(batch_size))
	{
		int curr_batch_size = pg_db.Tuples();
		if (!curr_batch_size)
			break;
		for (int i = 0; i < curr_batch_size; ++i, index += 2)
		{
			// save accession codes
			vec_acc_code_indices.push_back(t_acc_code_index(pg_db.getvalue(i, 4), index));
			vec_acc_code_indices.push_back(t_acc_code_index(pg_db.getvalue(i, 5), index+1));
			
			// save offset to strings
			pair_data.push_back(t_pair_data(index,
											index + 1,
											pg_db.getvalue(i, 0),	// gene_id1
											pg_db.getvalue(i, 1),	// gene_id2
											pg_db.getvalue(i, 2),	// pi1
											pg_db.getvalue(i, 3),	// pi2
											pg_db.getvalue(i, 6),	// gap1
											pg_db.getvalue(i, 7),	// gap2
																	// aln_begin
											lexical_cast<unsigned>(pg_db.getvalue(i, 8)),	
																	// aln_len
											lexical_cast<unsigned>(pg_db.getvalue(i, 9))
											));
			
		}								
	}
	
	// close cursor
	if (!pg_db.Close())
		return pg_print_sql_err("close the cursor to retrieve batch " + batch_num,
								pg_db, "CLOSE CURSOR",
								cerr);
	
	// nothing retrieved
	if (!vec_acc_code_indices.size())
	{
		cerr << "No matching entries in the alignment table could be found\n";
		return true;
	}
	
	
	//
	//	make acc_codes and PIs unique
	//
	construct_unique_acc_code_list(vec_acc_code_indices, pair_data, acc_codes);

	return true;
}
	






//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		Main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main(int argc, char* argv[])
{
	debug_func_cerr(5);
	
	// don't use C streams
	std::ios::sync_with_stdio(false);

	// get command line arguments
	t_program_args args;
	if (!GetArguments(argc, argv, args))
		return 1;

	
	//____________________________________________________________________________________
	// 		Connect to PostgreSQL
	//____________________________________________________________________________________

	t_pg_cursor pg_db(args.db_login.c_str(), "panda_cursor");
	if ( pg_db.ConnectionBad() )
	{
		cerr	<< "\nERROR!! Connection to database ["
							<< args.db_login << "] failed.\n"
							<< pg_db.ErrorMessage() << "\n\n";
		return 1;
	}
	pg_db.Commit();
	string PID = get_backend_pid(pg_db);
	args.open_log_files(PID);
	

	t_progress_indicator dots(args.VERBOSE(), (args.codeml ? 5 : 20));
	unsigned cnt_batches = 0;
	t_portable_timer timer;
	
	cerr	<< "\tPAML from " << args.paml << "\n";
	
	
	// use kaks_tmp as temp directory if not specified
	string tmp_file_path = args.tempdir.length() ?
									args.tempdir :
									path(argv[0]).branch_path().string() + "/kaks_tmp";
	
	str_search_and_replace(tmp_file_path, "%PID%", PID);

	// save original directory for restore later
	t_auto_restore_dir save_directory;
	bool remove_tmp_file = false;
	if(!prepare_temp_directory(tmp_file_path,
								false, 0.0, /*fixed_kappa, args.kappa,*/
								args.codeml, remove_tmp_file))
		return false;
	
	try
	{
	
	while (1)
	{
		//________________________________________________________________________________
		// 		get next batch number
		//________________________________________________________________________________
		t_batch_num batch_num (pg_db);

		// if last batch, try again after freeing hanging jobs
		if (!batch_num.get_next(args.job_category))
			break;
		
		dots.print_indicator(batch_num.text());

		pg_db.BeginTransaction();
		
		
		//________________________________________________________________________________
		// 		get pair data
		//________________________________________________________________________________
		t_vec_pair_data pair_data;
		vector<string> acc_codes;
		if (!get_sequence_pair_data(pg_db, args, batch_num.text(), pair_data, acc_codes))
			return 1;
		{	
			ostringstream is; is << " " << pair_data.size() << " pairs:";
			dots.print_indicator(is.str());
		}

		
		//________________________________________________________________________________
		// 		get cdna_sequences
		//________________________________________________________________________________
		vector<string> cdna_sequences;
		try
		{
		if (!get_cdna_sequences(pg_db, acc_codes, cdna_sequences, batch_num.text(), args))
			return 1;
		}
		catch(std::logic_error& e)
		{
			cerr << e.what() << std::endl;
		}
		catch(...)
		{
		}
		//________________________________________________________________________________
		// 		calculate kaks
		//________________________________________________________________________________
		// change to temp directory
		if (chdir(tmp_file_path.c_str()))
		{
            char buf[256];
			buf[255] = '\0';
			strerror_r(errno, buf, 255);
			cerr << "\nError!\n Could not change to " << tmp_file_path << "[errno="
				<< errno << "]\n"
					<< buf << "\n";
			return 1;
		}
		// cerr << boost::filesystem::current_path().string() << "\n";

		// calculate kaks
		vector<t_codeml_results> results;
		get_kaks_results(pair_data, acc_codes, cdna_sequences, results, dots, args);

		{	
			ostringstream is; is << " " << pair_data.size() << " results.";
			dots.print_indicator(is.str());
		}
		
		// write results
		if (!write_kaks_results(pg_db, pair_data, acc_codes, results, args))
			return 1;
		
		//________________________________________________________________________________
		// 		batch succeed, free batch number
		//________________________________________________________________________________
		batch_num.commit();
		if (!pg_db.Commit())
		{
			cerr	<<	"\n"
						"ERROR!! PostgreSQL error.\n\t"
						"Could not Commit the current transaction.\n\t"
					<<		pg_db.ErrorMessage()
					<<	std::endl;
			return 1;
		}
		++cnt_batches;
	}
	}
	catch(...)
	{
		return 1;
	}
	
	dots.finish();
	args.VERBOSE() << "All batches finished" << std::endl;
	args.VERBOSE() << timer.pretty_reset() << " for " << cnt_batches << " batches\n";


	return 0;
}

